import torch
import torchvision
from PIL import Image
from model import *

''''
模型验证步骤：
1. 加载图片并进行预处理：从指定路径加载图片，使用torchvision.transforms将其调整为32x32大小并转换为张量。
2. 加载模型：通过torch.load加载训练好的模型，并设置为评估模式。
3. 推理过程：将图片张量 reshape 为模型输入格式，关闭梯度计算后进行前向传播，输出预测结果及最大值索引
'''
img_path = '../../../imgs/dog.png'
img = Image.open(img_path)

transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32,32)),
                                            torchvision.transforms.ToTensor()])

image = transform(img)
print(image.shape)


model = torch.load("model_0.pth",map_location="cpu")
print(model)

image = torch.reshape(image, (1,3,32,32))
model.eval()
with torch.no_grad():
    output = model(image)
print(output)
print(torch.argmax(output))